import pandas
import numpy
from copy import deepcopy


class CustomIndicator(object):
    exclude_columns = ('floor0', 'floor1', 'upper0', 'upper1')

    def __init__(self, dataframe):
        self.dataframe = dataframe
        self.ma5 = self.dataframe.close.rolling(5, 0).mean()
        self.ma10 = self.dataframe.close.rolling(10, 0).mean()
        self.ma20 = self.dataframe.close.rolling(20, 0).mean()
        self.macd = self.macd()

    def candle(self):
        data = self.dataframe.close - self.dataframe.open
        data = data.mask(data < 0, -1).mask(data > 0, 1)
        data = data * (data.groupby((data != data.shift(1)).cumsum()).cumcount() + 1)
        return data

    def change(self):
        idx1 = self.dataframe.close.shift(1)
        data = (self.dataframe.close - idx1).fillna(0)
        return data / idx1 * 100

    def macd(self):
        ma12 = self.dataframe.close.ewm(span=12, min_periods=1, adjust=False).mean()
        ma26 = self.dataframe.close.ewm(span=26, min_periods=1, adjust=False).mean()
        data = ma12 - ma26
        data = data - data.ewm(span=9, min_periods=0, adjust=False).mean()
        return data

    def cloud(self):
        idx1 = (self.dataframe.close - self.dataframe.open).fillna(0)
        idx1 = idx1.mask(idx1 > 0, 1).mask(idx1 < 0, -1)
        idx2 = self.dataframe.close.where(
            (self.dataframe.close > self.dataframe.close.shift(1)) &
            (self.dataframe.close > self.dataframe.close.shift(2)) &
            (self.dataframe.close > self.dataframe.close.shift(3)) &
            (self.dataframe.close > self.dataframe.close.shift(4)),
            numpy.inf
        )
        idx2 = idx2.where(numpy.isinf(idx2), 1)
        tmp2 = self.dataframe.close.where(
            (self.dataframe.close < self.dataframe.close.shift(1)) &
            (self.dataframe.close < self.dataframe.close.shift(2)) &
            (self.dataframe.close < self.dataframe.close.shift(3)) &
            (self.dataframe.close < self.dataframe.close.shift(4)),
            numpy.inf
        )
        tmp2 = tmp2.where(numpy.isinf(tmp2), -1)
        idx2 = idx2.where(idx2 != numpy.inf, tmp2)
        idx2 = idx2.where(idx2 != numpy.inf, 0)
        idx3 = (self.dataframe.close - self.ma20).fillna(0)
        idx3 = idx3.mask(idx3 > 0, 1).mask(idx3 < 0, -1)
        idx4 = (self.ma5 - self.ma20).fillna(0)
        idx4 = idx4.mask(idx4 > 0, 1).mask(idx4 < 0, -1)
        idx5 = self.macd.where(
            (self.macd > 0) &
            (self.macd.shift(1) > 0) &
            (self.macd.shift(2) > 0),
            numpy.inf
        )
        idx5 = idx5.where(numpy.isinf(idx5), 1)
        tmp5 = self.macd.where(
            (self.macd < 0) &
            (self.macd.shift(1) < 0) &
            (self.macd.shift(2) < 0),
            numpy.inf
        )
        tmp5 = tmp5.where(numpy.isinf(tmp5), -1)
        idx5 = idx5.where(idx5 != numpy.inf, tmp5)
        idx5 = idx5.where(idx5 != numpy.inf, 0)
        return idx1 + idx2 + idx3 + idx4 + idx5

    def sun(self):
        macd = self.macd.mask(self.macd > 0, 1).mask(self.macd < 0, -1)
        macd = macd * (macd.groupby((macd != macd.shift(1)).cumsum()).cumcount() + 1)
        return macd.fillna(0)

    def deviate(self):
        data = pandas.Series(data=numpy.nan, index=self.dataframe.index)
        macd = self.dataframe.sun
        hm, hp, lm, lp = list(), list(), list(), list()
        for index, value in macd.items():
            if index > 0:
                grade_val = macd[index - 1]
                grade_abs = abs(grade_val)
                if grade_abs > abs(value):
                    start = int(index - grade_abs)
                    if grade_val > 0:
                        hm.append(self.macd[start: index].max())
                        hp.append(self.dataframe.high[start: index].max())
                    else:
                        lm.append(self.macd[start: index].min())
                        lp.append(self.dataframe.low[start: index].min())
                if value == 0:
                    data[index] = 0
                elif value == -1:
                    if len(hm) > 1 and hm[-1] < hm[-2] and hp[-1] >= hp[-2]:
                        data[index] = -1
                    else:
                        data[index] = 0
                elif value == 1:
                    if len(lm) > 1 and lm[-1] > lm[-2] and lp[-1] <= lp[-2]:
                        data[index] = 1
                    else:
                        data[index] = 0
        data = data.ffill()
        return data.fillna(0)

    def fly(self):
        data = self.dataframe.close.shift(1)
        data = (self.ma20.shift(1) - data).fillna(0)
        return data.mask(data > 0, 1).mask(data < 0, 0)

    def moon(self):
        idx1 = (self.ma5 - self.ma10).fillna(0)
        idx1 = idx1.mask(idx1 > 0, 1).mask(idx1 < 0, -1)
        idx2 = (self.ma10 - self.ma20).fillna(0)
        idx2 = idx2.mask(idx2 > 0, 1).mask(idx2 < 0, -1)
        idx3 = (self.ma10 - self.ma10.shift(1)).fillna(0)
        idx3 = idx3.mask(idx3 > 0, 1).mask(idx3 < 0, -1)
        idx4 = (self.ma20 - self.ma20.shift(1)).fillna(0)
        idx4 = idx4.mask(idx4 > 0, 1).mask(idx4 < 0, -1)
        idx5 = (self.ma10 - self.ma20).fillna(0)
        idx5 = idx5 - (self.ma10.shift(1) - self.ma20.shift(1)).fillna(0)
        idx5 = idx5.mask(idx5 > 0, 1).mask(idx5 < 0, -1)
        idx6 = self.macd.where(
            (self.macd > 0) &
            (self.macd.shift(1) > 0) &
            (self.macd.shift(2) > 0),
            numpy.inf
        )
        idx6 = idx6.where(numpy.isinf(idx6), 1)
        tmp6 = self.macd.where(
            (self.macd < 0) &
            (self.macd.shift(1) < 0) &
            (self.macd.shift(2) < 0),
            numpy.inf
        )
        tmp6 = tmp6.where(numpy.isinf(tmp6), -1)
        idx6 = idx6.where(idx6 != numpy.inf, tmp6)
        idx6 = idx6.where(idx6 != numpy.inf, 0)
        return idx1 + idx2 + idx3 + idx4 + idx5 + idx6

    def road(self):
        idx1 = (self.dataframe.close - self.ma20).fillna(0)
        idx1 = idx1.mask(idx1 > 0, 1).mask(idx1 < 0, -1)
        idx2 = (self.ma20 - self.ma20.shift(1)).fillna(0)
        idx2 = idx2.mask(idx2 > 0, 1).mask(idx2 < 0, -1)
        temp = (self.ma20.shift(1) - self.dataframe.close.shift(1)).fillna(0)
        temp = temp.mask(temp > 0, 1).mask(temp < 0, -1)
        temp = idx1 + temp
        idx3 = temp.where(temp == 2, numpy.inf)
        idx3 = idx3.where(numpy.isinf(idx3), 1)
        tmp3 = temp.where(temp == -2, numpy.inf)
        tmp3 = tmp3.where(numpy.isinf(tmp3), -1)
        idx3 = idx3.where(idx3 != numpy.inf, tmp3)
        idx3 = idx3.where(idx3 != numpy.inf, 0)
        return idx1 + idx2 + idx3

    def cma10(self):
        data = self.dataframe.close.rolling(10, 0).mean()
        idx1 = (data - data.shift(1)).fillna(0)
        idx1 = idx1.mask(idx1 > 0, 2).mask(idx1 < 0, -2)
        idx2 = (self.dataframe.close - data).fillna(0)
        idx2 = idx2.mask(idx2 > 0, 2).mask(idx2 < 0, -2)
        return idx1 + idx2

    def pma20(self):
        data = self.dataframe.close.rolling(20, 0).mean()
        data = (data - data.shift(1)).fillna(0)
        return data.mask(data > 0, 2).mask(data < 0, -2)

    def stoch(self):
        floor = self.dataframe.low.rolling(9, 0).min()
        floor.fillna(value=self.dataframe.low.expanding().min(), inplace=True)
        upper = self.dataframe.high.rolling(9, 0).max()
        upper.fillna(value=self.dataframe.high.expanding().max(), inplace=True)
        datas = (self.dataframe.close - floor) / (upper - floor) * 100
        datas = datas.fillna(value=100)
        datas = datas.ewm(com=2, adjust=False).mean()
        datas = (datas - datas.shift(1)).fillna(0)
        datas = datas.mask(datas < 0, -2)
        datas = datas.mask(datas > 0, 2)
        return datas

    def state(self):
        data = self.dataframe.total
        data = data.mask(data > 0, 1).mask(data < 0, -1).mask(data == 0, numpy.nan).ffill()
        data = data * (data.groupby((data != data.shift(1)).cumsum()).cumcount() + 1)
        return data.fillna(0)

    def grade(self):
        data1 = pandas.Series(data=numpy.nan, index=self.dataframe.index)
        data2 = pandas.Series(data=numpy.nan, index=self.dataframe.index)
        p_idx = list()
        for index, value in self.dataframe.state.items():
            if index > 1:
                if value < 0 < self.dataframe.state[index - 1] and self.dataframe.state[index - 2] > 0:
                    if len(p_idx) > 0 and p_idx[-1][1] == -1:
                        p_idx[-1] = [index, -1]
                    else:
                        p_idx.append([index, -1])
                if self.dataframe.state[index - 1] < 0 < value and self.dataframe.state[index - 2] < 0:
                    if len(p_idx) > 0 and p_idx[-1][1] == 1:
                        p_idx[-1] = [index, 1]
                    else:
                        p_idx.append([index, 1])
            if value in (-1, 1):
                if len(p_idx) > 1:
                    _p_idx = deepcopy(p_idx)
                    if _p_idx[-1][0] != index:
                        if _p_idx[-1][1] == value:
                            _p_idx[-1] = [index, value]
                        else:
                            _p_idx.append([index, value])
                    if len(_p_idx) > 1:
                        data1[index] = _p_idx[-1][0] - _p_idx[-2][0]
                    if len(_p_idx) > 3:
                        data2[index] = _p_idx[-3][0] - _p_idx[-4][0]
        data1 = data1.ffill()
        data2 = data2.ffill()
        return data1.fillna(0), data2.fillna(0)

    def floor0(self):
        data = pandas.Series(data=numpy.nan, index=self.dataframe.index)
        temp = self.dataframe.state[self.dataframe.state == 1]
        for index in temp.index:
            fi = index - self.dataframe.iloc[index].grade
            if fi >= 0:
                data[index] = self.dataframe.low[int(fi):index].min()
        return data.ffill()

    def upper0(self):
        data = pandas.Series(data=numpy.nan, index=self.dataframe.index)
        temp = self.dataframe.state[self.dataframe.state == -1]
        for index in temp.index:
            fi = index - self.dataframe.iloc[index].grade
            if fi >= 0:
                data[index] = self.dataframe.high[int(fi):index].max()
        return data.ffill()

    def floor1(self):
        data = pandas.Series(data=numpy.nan, index=self.dataframe.index)
        temp = self.dataframe.loc[self.dataframe.state.eq(1)]
        maps = dict()
        for index, row in temp.iterrows():
            if row.extent in maps:
                data[index] = maps[row.extent][-1]
            if row.grade not in maps:
                maps[row.grade] = [row.floor0]
            else:
                maps[row.grade].append(row.floor0)
        return data.ffill()

    def upper1(self):
        data = pandas.Series(data=numpy.nan, index=self.dataframe.index)
        temp = self.dataframe.loc[self.dataframe.state.eq(-1)]
        maps = dict()
        for index, row in temp.iterrows():
            if row.extent in maps:
                data[index] = maps[row.extent][-1]
            if row.grade not in maps:
                maps[row.grade] = [row.upper0]
            else:
                maps[row.grade].append(row.upper0)
        return data.ffill()

    def undulate(self):
        data = pandas.Series(data=numpy.nan, index=self.dataframe.index)
        temp = self.dataframe.loc[self.dataframe.state.eq(1)]
        temp = (temp.floor0 / temp.upper0 - 1) * 100
        data = data.fillna(temp)
        temp = self.dataframe.loc[self.dataframe.state.eq(-1)]
        temp = (temp.upper0 / temp.floor0 - 1) * 100
        data = data.fillna(temp)
        return data.ffill()

    def surplus(self):
        data = pandas.Series(data=numpy.nan, index=self.dataframe.index)
        temp = self.dataframe.loc[self.dataframe.state.eq(1)]
        temp = (temp.floor0 - temp.floor1) / temp.floor1 * 100
        data = data.fillna(temp)
        temp = self.dataframe.loc[self.dataframe.state.eq(-1)]
        temp = (temp.upper0 - temp.upper1) / temp.upper1 * 100
        data = data.fillna(temp)
        return data.ffill()

    def deficit(self):
        data = pandas.Series(data=numpy.nan, index=self.dataframe.index)
        temp = self.dataframe.loc[self.dataframe.state.eq(1)]
        temp = (temp.upper0 - temp.upper1) / temp.upper1 * 100
        data = data.fillna(temp)
        temp = self.dataframe.loc[self.dataframe.state.eq(-1)]
        temp = (temp.floor0 - temp.floor1) / temp.floor1 * 100
        data = data.fillna(temp)
        return data.ffill()

    def profit(self):
        data = pandas.Series(data=numpy.nan, index=self.dataframe.index)
        temp = self.dataframe.loc[self.dataframe.state.eq(1)]
        temp = (temp.upper0 / temp.close - 1) * 100
        data = data.fillna(temp)
        temp = self.dataframe.loc[self.dataframe.state.eq(-1)]
        temp = (temp.close - temp.floor0) / temp.close * 100
        data = data.fillna(temp)
        return data.ffill()

    def trend(self):
        values = pandas.Series(numpy.nan, self.dataframe.index)
        for index, row in self.dataframe.iterrows():
            value, point = 0, int(index - abs(row.state))
            if row.state and point >= 0:
                close = self.dataframe.iloc[point].close
                value = (row.close - close) / close * 100
            values[index] = value
        return pandas.Series(values, self.dataframe.index)

    def result(self):
        self.dataframe['candle'] = self.candle()
        self.dataframe['change'] = self.change()
        self.dataframe['cloud'] = self.cloud()
        self.dataframe['sun'] = self.sun()
        self.dataframe['deviate'] = self.deviate()
        self.dataframe['fly'] = self.fly()
        self.dataframe['moon'] = self.moon()
        self.dataframe['road'] = self.road()
        self.dataframe['total'] = self.cma10() + self.stoch() + self.pma20()
        self.dataframe['state'] = self.state()
        self.dataframe['grade'], self.dataframe['extent'] = self.grade()
        self.dataframe['floor0'] = self.floor0()
        self.dataframe['floor1'] = self.floor1()
        self.dataframe['upper0'] = self.upper0()
        self.dataframe['upper1'] = self.upper1()
        self.dataframe['undulate'] = self.undulate()
        self.dataframe['surplus'] = self.surplus()
        self.dataframe['deficit'] = self.deficit()
        self.dataframe['profit'] = self.profit()
        self.dataframe['trend'] = self.trend()
        self.dataframe = self.dataframe[[column for column in self.dataframe.columns if column not in self.exclude_columns]]
        return self.dataframe.fillna(0)
